Greptile Summary

This PR fixes the distributed JAX MoE auxiliary-loss backward pass by removing an incorrect g.reshape(1) call that was converting the scalar cotangent (g, rank-0) into a shape-(1,) tensor (rank-1). Under Shardy, the FusedMoEAuxLossBwdPrimitive.shardy_sharding_rule declares the grad_aux_loss operand as rank-0, so the reshape caused a rank-mismatch failure during distributed lowering; the FFI C++ backend already accepts a scalar buffer natively.

• Root cause removed: g.reshape(1) in _fused_moe_aux_loss_bwd is deleted; the raw scalar cotangent g is now passed directly to fused_moe_aux_loss_bwd, matching the rank-0 declaration in the Shardy rule.
• No API surface changed: the FusedMoEAuxLossBwdPrimitive.abstract method only reads grad_aux_loss_aval.dtype and is shape-agnostic, so the abstract evaluation is unaffected by the fix.

Confidence Score: 5/5

The change is a one-line deletion that aligns the cotangent shape with the existing Shardy rule and FFI contract; no new behaviour is introduced and the distributed backward test already covers this path.

The fix removes a single reshape that was causing rank mismatches under Shardy distributed lowering. The abstract evaluation in the primitive is shape-agnostic (reads only dtype), the FFI backend already accepts a scalar, and the Shardy rule explicitly documents the rank-0 expectation with an explanatory comment.

No files require special attention; the single changed hunk in transformer_engine/jax/router.py is self-contained and consistent with the primitive's declared sharding contract.

Important Files Changed

Sequence Diagram

sequenceDiagram
    participant JAX as JAX Autograd
    participant BWD as _fused_moe_aux_loss_bwd
    participant API as fused_moe_aux_loss_bwd
    participant PRIM as FusedMoEAuxLossBwdPrimitive
    participant FFI as C++ FFI kernel

    JAX->>BWD: g (rank-0 scalar cotangent)
    Note over BWD: BEFORE: grad_aux_loss = g.reshape(1) shape (1,) rank-1, Shardy expects rank-0
    Note over BWD: AFTER: pass g directly shape () rank-0, Shardy rule satisfied
    BWD->>API: fused_moe_aux_loss_bwd(const_buf, tokens_per_expert, g, num_tokens)
    API->>PRIM: outer_primitive.bind(...)
    PRIM->>PRIM: shardy_sharding_rule declares grad_aux_loss as rank-0
    PRIM->>FFI: ffi.ffi_lowering scalar buffer accepted natively
    FFI-->>JAX: grad_probs [num_tokens, num_experts]

_{Reviews (1): Last reviewed commit: "[JAX] Keep MoE aux-loss cotangent scalar" | Re-trigger Greptile}

/te-ci jax

/te-ci L1 jax

/te-ci L2 jax

THE H100 takes forever to queue in this pipeline: https://gitlab-master.nvidia.com/dl/transformerengine/transformerengine/-/pipelines/59294969 but the B200 and GB200 jobs all passed.

H100 also passed in previous jobs. Merging